from rag_test.pingan.es_insert import insert_into_es
from rag_test.pingan.pingan_milvus_storage import init_milvus, insert_into_milvus
from rag_test.pingan.pingan_pdf_loader import read_pdf
from rag_test.pingan.pingan_text_embedding import text_embeddings
from rag_test.pingan.pingan_text_splitter import split_text

if __name__ == '__main__':
    # 替换成你的PDF文件路径
    pdf_path = "平安保险用户手册.pdf"
    sample_pdf_contents = read_pdf(pdf_path)
    # 调用切分函数
    sample_chunks = split_text(sample_pdf_contents, chunk_size=100, chunk_overlap=20)

    # 打印结果
    # print(f"切分后得到{len(sample_chunks)}个片段：")
    # for i, chunk in enumerate(sample_chunks, 1):
    #     print(f"\n片段{i}（第{sample_chunks['page_num']}页）：")
    #     print(chunk['chunk'])

    # 提取文本内容（只需要chunk部分）
    #texts = [chunk["chunk"] for chunk in sample_chunks]

    # 调用向量化函数
    #vectors = text_embeddings(texts)

    # # 打印结果
    # print(f"生成了{len(vectors)}个向量")
    # # 打印第一个向量的前5个数值（向量很长，只看一部分）
    # if vectors:
    #     print("第一个向量的前5个值：", vectors[0][:5])
    #     print("向量长度（维度）：", len(vectors[0]))  # 应该是1024
  # 1. 初始化Milvus集合
    #collection = init_milvus()

    # # 2. 模拟数据（实际使用时从前面的步骤获取）
    # sample_chunks = [
    #     {"page_num": 1, "chunk": "平安保险成立于1988年，总部位于深圳。"},
    #     {"page_num": 1, "chunk": "旗下有多个保险产品，包括寿险、车险、健康险等。"},
    #     {"page_num": 2, "chunk": "平安福是公司的明星寿险产品，2024年推出了新版。"}
    # ]

    # 模拟向量（实际使用时调用text_embeddings函数生成）
    # 这里用随机数代替真实向量（仅测试用）
    #sample_embeddings = [np.random.rand(1024).tolist() for _ in range(3)]

    # 3. 插入数据到Milvus
    #insert_into_milvus(collection, sample_chunks, vectors)

    insert_into_es(sample_chunks)
